import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import os
from matplotlib.ticker import FuncFormatter
import matplotlib.colors as mcolors

# 创建目录保存图片（如果不存在）
output_dir = "benchmark_analysis_plots"
os.makedirs(output_dir, exist_ok=True)

# 您提供的原始数据
raw_data = [
    [6, 297856, 3308650, 22068224],
    [324656, 321775425, 11221288859, 78524397760],
    [6378048, 416959760, 1070392228, 5438720000],
    [260580, 49182512, 2439293716, 54550132311],
    [4334080, 102914560, 355912387, 1695767040],
    [50336, 29077952, 482594905, 11618640000],
    [80480, 14217117, 276199106, 3752390356],
    [18176675, 3909825077, 11952352871, 52202026615],
    [6, 27158732, 1270964114, 78524397760]
]

# Case名称映射
case_names = [
    "Microbenchmarks", 
    "Rodinia", 
    "Polybench", 
    "Industry HPC", 
    "CUBLAS", 
    "CUTLASS", 
    "Deepbench", 
    "MLPerf", 
    "All"
]

# Scale名称
scale_names = ["Min", "Medium", "Average", "Max"]

# 性能目标 (KIPS) - 用于热力图批量生成
kips_targets_heatmaps = [2, 20, 40, 100, 200, 400, 800, 1000, 2000, 4000, 16000, 25000, 50000, 100000, 200000]

# 转换数据为NumPy数组以便计算
data = np.array(raw_data)
num_cases, num_scales = data.shape

# 初始化一个字典来存储结果
results = {}

for kips in kips_targets_heatmaps:
    # 计算当前KIPS目标下的执行时间 (秒)
    time_matrix = data / (kips * 1000)
    results[kips] = time_matrix

# 时间单位转换函数
def format_time(seconds, pos=None):
    if seconds >= 86400:  # 天
        return f'{seconds/86400:.1f}d'
    elif seconds >= 3600:  # 小时
        return f'{seconds/3600:.1f}h'
    elif seconds >= 60:    # 分钟
        return f'{seconds/60:.1f}m'
    else:                  # 秒
        return f'{seconds:.1f}s'

# 设置全局绘图样式
plt.rcParams['figure.figsize'] = (14, 10)
plt.rcParams['font.size'] = 10
plt.rcParams['axes.grid'] = True
plt.rcParams['font.family'] = 'DejaVu Sans'

# 图表 1：固定规模（规模4），分析不同测试用例
scale_idx = 3  # 第四个规模
max_time_seconds = 5 * 86400  # 5天的秒数

plt.figure(figsize=(16, 12))
for case_idx in range(num_cases):
    times = [results[kips][case_idx, scale_idx] for kips in kips_targets_heatmaps]
    
    # 区分小于等于5天和大于5天的数据点
    valid_indices = [i for i, t in enumerate(times) if t <= max_time_seconds]
    invalid_indices = [i for i, t in enumerate(times) if t > max_time_seconds]
    
    valid_kips = [kips_targets_heatmaps[i] for i in valid_indices]
    valid_times = [times[i] for i in valid_indices]
    invalid_kips = [kips_targets_heatmaps[i] for i in invalid_indices]
    invalid_times = [times[i] for i in invalid_indices]
    
    # 绘制所有数据点连线
    plt.plot(kips_targets_heatmaps, times, marker='', linestyle='-', color='lightgray', linewidth=1, alpha=0.5)
    
    # 绘制小于等于5天的数据点（正常颜色）
    if valid_indices:
        line = plt.plot(valid_kips, valid_times, marker='o', markersize=8, linewidth=3, 
                       label=case_names[case_idx], alpha=0.8)
        
        # 标记最后一个有效点
        if len(valid_times) > 0:
            last_idx = len(valid_times) - 1
            plt.annotate(case_names[case_idx], 
                        xy=(valid_kips[last_idx], valid_times[last_idx]),
                        xytext=(10, 0), textcoords='offset points',
                        fontsize=9, fontweight='bold',
                        bbox=dict(boxstyle="round,pad=0.3", fc=line[0].get_color(), alpha=0.7))
    
    # 绘制大于5天的数据点（红色）
    if invalid_indices:
        plt.plot(invalid_kips, invalid_times, marker='x', markersize=10, linewidth=3, 
                color='red', alpha=0.7, markeredgewidth=2)

plt.xscale('log')
plt.yscale('log')
plt.xlabel('Performance (KIPS)', fontsize=14, fontweight='bold')
plt.ylabel('Execution Time', fontsize=14, fontweight='bold')
plt.title(f'Execution Time vs. Performance for Different Benchmarks (Scale: {scale_names[scale_idx]})\nRed X: Execution Time > 5 Days', 
          fontsize=16, fontweight='bold', pad=20)

# 使用自定义的时间格式化器
plt.gca().yaxis.set_major_formatter(FuncFormatter(format_time))

# 添加5天的时间参考线
plt.axhline(y=max_time_seconds, color='red', linestyle='--', linewidth=2, alpha=0.7, 
           label='5 Days Threshold')
plt.text(kips_targets_heatmaps[0] * 1.1, max_time_seconds * 1.1, '5 Days', 
         fontsize=12, fontweight='bold', color='red', va='bottom')

plt.grid(True, which='both', linestyle='--', linewidth=0.5, alpha=0.7)
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=11)
plt.tight_layout()
plt.savefig(f'{output_dir}/chart1_benchmarks_at_scale_{scale_names[scale_idx].lower()}.png', dpi=300, bbox_inches='tight')
plt.savefig(f'{output_dir}/chart1_benchmarks_at_scale_{scale_names[scale_idx].lower()}.pdf', bbox_inches='tight')
plt.close()

# 图表 2：固定测试用例（MLPerf），分析不同规模
case_idx = 7

plt.figure(figsize=(16, 12))
for scale_idx in range(num_scales):
    times = [results[kips][case_idx, scale_idx] for kips in kips_targets_heatmaps]
    
    # 区分小于等于5天和大于5天的数据点
    valid_indices = [i for i, t in enumerate(times) if t <= max_time_seconds]
    invalid_indices = [i for i, t in enumerate(times) if t > max_time_seconds]
    
    valid_kips = [kips_targets_heatmaps[i] for i in valid_indices]
    valid_times = [times[i] for i in valid_indices]
    invalid_kips = [kips_targets_heatmaps[i] for i in invalid_indices]
    invalid_times = [times[i] for i in invalid_indices]
    
    # 绘制所有数据点连线
    plt.plot(kips_targets_heatmaps, times, marker='', linestyle='-', color='lightgray', linewidth=1, alpha=0.5)
    
    # 绘制小于等于5天的数据点
    if valid_indices:
        line = plt.plot(valid_kips, valid_times, marker='s', markersize=8, linewidth=3, 
                       label=f'{scale_names[scale_idx]} Scale', alpha=0.8)
        
        # 标记最后一个有效点
        if len(valid_times) > 0:
            last_idx = len(valid_times) - 1
            plt.annotate(scale_names[scale_idx], 
                        xy=(valid_kips[last_idx], valid_times[last_idx]),
                        xytext=(10, 0), textcoords='offset points',
                        fontsize=9, fontweight='bold',
                        bbox=dict(boxstyle="round,pad=0.3", fc=line[0].get_color(), alpha=0.7))
    
    # 绘制大于5天的数据点（红色）
    if invalid_indices:
        plt.plot(invalid_kips, invalid_times, marker='x', markersize=10, linewidth=3, 
                color='red', alpha=0.7, markeredgewidth=2)

plt.xscale('log')
plt.yscale('log')
plt.xlabel('Performance (KIPS)', fontsize=14, fontweight='bold')
plt.ylabel('Execution Time', fontsize=14, fontweight='bold')
plt.title(f'Execution Time vs. Performance for Different Scales ({case_names[case_idx]})\nRed X: Execution Time > 5 Days', 
          fontsize=16, fontweight='bold', pad=20)

# 使用自定义的时间格式化器
plt.gca().yaxis.set_major_formatter(FuncFormatter(format_time))

# 添加5天的时间参考线
plt.axhline(y=max_time_seconds, color='red', linestyle='--', linewidth=2, alpha=0.7, 
           label='5 Days Threshold')
plt.text(kips_targets_heatmaps[0] * 1.1, max_time_seconds * 1.1, '5 Days', 
         fontsize=12, fontweight='bold', color='red', va='bottom')

plt.grid(True, which='both', linestyle='--', linewidth=0.5, alpha=0.7)
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=11)
plt.tight_layout()
plt.savefig(f'{output_dir}/chart2_scales_for_{case_names[case_idx].lower()}.png', dpi=300, bbox_inches='tight')
plt.savefig(f'{output_dir}/chart2_scales_for_{case_names[case_idx].lower()}.pdf', bbox_inches='tight')
plt.close()

# 图表 3：批量生成热力图 - 为每个性能目标生成一个热力图
max_time_seconds = 5 * 86400  # 5天的秒数

# 创建自定义颜色映射：绿色到黄色到红色，并在5天处设置明确的界限
colors = ['#00ff00', '#80ff00', '#ffff00', '#ff8000', '#ff0000', '#800000']
cmap = mcolors.LinearSegmentedColormap.from_list('custom_red_yellow_green', colors, N=100)

for target_kips in kips_targets_heatmaps:
    time_matrix = results[target_kips]
    
    plt.figure(figsize=(16, 12))
    
    # 使用对数色标，设置vmax为5天的对数值
    log_time = np.log10(time_matrix)
    im = plt.imshow(log_time, cmap=cmap, aspect='auto', 
                   vmin=0, vmax=np.log10(max_time_seconds))

    # 设置颜色条
    cbar = plt.colorbar(im, label='Execution Time', shrink=0.8)
    cbar.ax.tick_params(labelsize=10)
    
    # 在颜色条上添加时间单位标注
    cbar_ticks = [0, 1, 2, 3, 4, np.log10(max_time_seconds)]
    cbar_tick_labels = ['1s', '10s', '100s', '1000s', '10000s', '5d']
    cbar.set_ticks(cbar_ticks)
    cbar.set_ticklabels(cbar_tick_labels)

    # 设置坐标轴标签
    plt.xticks(range(num_scales), scale_names, fontsize=12, fontweight='bold')
    plt.yticks(range(num_cases), case_names, fontsize=12, fontweight='bold')
    plt.xlabel('Problem Scale', fontsize=14, fontweight='bold')
    plt.ylabel('Benchmark Suite', fontsize=14, fontweight='bold')
    target_ips = target_kips
    ips_unit = "KIPS"
    target_mips = target_kips / 1000
    if (target_mips >= 1):
        target_ips = target_mips
        ips_unit = "MIPS"
    plt.title(f'Execution Time at {target_ips:,} {ips_unit}\n(Green: Fast, Yellow: Medium, Red: Slow, >5d: Maximum Red)', 
              fontsize=16, fontweight='bold', pad=20)

    # 在热力图上显示所有数值
    for i in range(num_cases):
        for j in range(num_scales):
            time_val = time_matrix[i, j]
            
            # 根据时间值选择合适的显示格式
            if time_val > 86400:  # 大于1天
                display_text = f'{time_val/86400:.1f}d'
            elif time_val > 3600:  # 大于1小时
                display_text = f'{time_val/3600:.1f}h'
            elif time_val > 60:  # 大于1分钟
                display_text = f'{time_val/60:.1f}m'
            else:  # 显示秒
                display_text = f'{time_val:.1f}s'
            
            # 根据背景色亮度调整文字颜色以确保可读性
            bg_brightness = log_time[i, j] / np.log10(max_time_seconds)
            if bg_brightness > 0.6:  # 背景较亮时使用深色文字
                font_color = "black"
            else:
                font_color = "white"
            
            # 如果超过5天，特别标记
            if time_val > max_time_seconds:
                display_text = f'>{display_text}'
                font_color = "white"
                plt.text(j, i, display_text, ha="center", va="center", 
                        color=font_color, fontsize=9, fontweight='bold',
                        bbox=dict(boxstyle="round,pad=0.2", fc="darkred", alpha=0.8))
            else:
                plt.text(j, i, display_text, ha="center", va="center", 
                        color=font_color, fontsize=9, fontweight='bold')

    plt.tight_layout()
    # 保存热力图，文件名包含KIPS值
    kips_str = f"{target_kips//1000}k" if target_kips >= 1000 else str(target_kips)
    plt.savefig(f'{output_dir}/heatmap_{kips_str}_kips.png', dpi=300, bbox_inches='tight')
    plt.savefig(f'{output_dir}/heatmap_{kips_str}_kips.pdf', bbox_inches='tight')
    plt.close()

# 创建并保存完整的数据表格
all_data = []
for target_kips in kips_targets_heatmaps:
    time_matrix = results[target_kips]
    for case_idx in range(num_cases):
        for scale_idx in range(num_scales):
            time_val = time_matrix[case_idx, scale_idx]
            all_data.append({
                'KIPS': target_kips,
                'Benchmark': case_names[case_idx],
                'Scale': scale_names[scale_idx],
                'Instructions': data[case_idx, scale_idx],
                'Time_seconds': time_val,
                'Time_formatted': format_time(time_val),
                'Within_5_days': time_val <= max_time_seconds
            })

all_df = pd.DataFrame(all_data)
all_df.to_csv(f'{output_dir}/all_execution_times_data.csv', index=False)

# 创建详细的分析报告
with open(f'{output_dir}/analysis_report.txt', 'w') as f:
    f.write("BENCHMARK EXECUTION TIME ANALYSIS REPORT\n")
    f.write("=" * 60 + "\n\n")
    f.write(f"Analysis Threshold: 5 Days ({max_time_seconds} seconds)\n\n")
    
    f.write("SUMMARY:\n")
    f.write("-" * 40 + "\n")
    
    total_configs = len(all_data)
    within_5_days = len(all_df[all_df['Within_5_days'] == True])
    
    f.write(f"Total configurations analyzed: {total_configs}\n")
    f.write(f"Configurations within 5 days: {within_5_days} ({within_5_days/total_configs*100:.1f}%)\n")
    f.write(f"Configurations exceeding 5 days: {total_configs - within_5_days} ({(total_configs - within_5_days)/total_configs*100:.1f}%)\n\n")
    
    f.write("CONFIGURATIONS EXCEEDING 5 DAYS:\n")
    f.write("-" * 40 + "\n")
    exceeding_configs = all_df[all_df['Within_5_days'] == False]
    if len(exceeding_configs) > 0:
        for _, row in exceeding_configs.iterrows():
            f.write(f"{row['Benchmark']} - {row['Scale']} at {row['KIPS']} KIPS: {row['Time_formatted']}\n")
    else:
        f.write("None - All configurations complete within 5 days!\n")

print(f"图表已保存到目录: {output_dir}/")
print("1. ✓ 区分显示：≤5天正常显示，>5天用红色X标记")
print("2. ✓ 热力图完整显示所有数据，>5天的用最深红色表示")
print("3. ✓ 热力图中>5天的")
print("\n生成的文件:")
print(f"- 15个热力图文件: heatmap_1k_kips.png/pdf 到 heatmap_200k_kips.png/pdf")
print(f"- 2个趋势图文件")
print(f"- 1个详细数据CSV文件")
print(f"- 1个分析报告")
